from langgraph.prebuilt import create_react_agent
from langchain_core.tools import tool
import random
from customize.get_ollama import GetOllama
from langchain_core.messages import SystemMessage, HumanMessage
from customize.save_image import save_graph

light_state = False


@tool
def people_detection() -> bool:
    """
    检测环境中是否有人，并返回检测结果。
    :return: 当检测到人是放回True，没有检测到返回False
    """
    value = random.choice([True, False])
    # value = False
    print(f"环境检测结果{value}")
    return value


@tool
def light_control(control_instruction: bool) -> bool:
    """
    通过输入灯光控制命令，控制灯的开和关。
    :param control_instruction:True表示开灯，False表示关灯。
    :return: 灯开则返回True，灯关闭返回False
    """

    if control_instruction:
        global light_state
        print("Turn on the light!")
        light_state = True
        print("开灯动作")
        return True
    else:
        print("Turn off the light!")
        light_state = False
        print("关灯动作")
        return False


@tool
def light_detection() -> bool:
    """
    检测当前灯光的状态
    :return: 开返回True，关返回False
    """
    print(f"灯光检测结果：{light_state}")
    return light_state


tools = [people_detection, light_control, light_detection]
llm = GetOllama(model_name="qwen2.5:14b", model_type=1, temperature=0)()
system = """
你作为一个管理助手，接收用户输入或工具的输入，并依据以下要求做出反应：
当收到用户输入的文本时，首先剖析该文本的任务意图，任务可能涵盖检测环境中是否有人、检测环境的开关灯状态以及控制灯光开关这三种。
请依照以下逻辑调用工具：
1. 若需检测环境是否开灯，调用工具 [light_detection] 进行检测。
2. 若需开灯/关灯，调用工具 [light_control] 进行操作。
3. 若需检测环境中是否有人，需先调用工具 [light_detection] 检测环境是否开灯，再依据检测结果选择以下一项进行操作：
   1）处于开灯状态时，使用工具 [people_detection] 检测环境中是否有人。
   2）处于关灯状态时，先调用工具 [light_control] 开灯，再使用工具 [people_detection] 检测环境中是否有人。
4. 其它要求，组合1~3的操作，规划好执行的步骤，根据规划一步一步调用工具执行任务，并把每一步执行的结果记录。
将每一步的操作都反馈给用户，并把最终的环境状况和开关灯结果告知用户。
"""
sys_msg = SystemMessage(content=system)
agent = create_react_agent(
    model=llm,
    tools=tools,
    state_modifier=sys_msg,
)
save_graph(agent, "react_agent.png")

inputs = {"messages": [("user", "1.帮我看看有没有人；2.如果有人，看看那个人叫什么名字？")]}
for s in agent.stream(inputs, stream_mode="values"):
    message = s["messages"][-1]
    if isinstance(message, tuple):
        print(message)
    else:
        message.pretty_print()
